package tk.hongbo.stream.coder;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * Created by HongboZhao on 2022/1/26.
 */
public class AvcEncoder {

    private int TIMEOUT_USEC = 12000;

    private MediaCodec mediaCodec;
    int m_width;
    int m_height;
    int m_framerate;

    public byte[] configbyte;

    private OnListener listener;

    private ExecutorService executor = Executors.newCachedThreadPool();

    public AvcEncoder(int width, int height, int framerate, OnListener onListener) {
        m_width = width;
        m_height = height;
        m_framerate = framerate;
        this.listener = onListener;

        MediaFormat mediaFormat = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_HEVC, width, height);
        mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar);
        mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, width * height * 5);
        mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, m_framerate);
        mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1);
        try {
            mediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_HEVC);
        } catch (IOException e) {
            e.printStackTrace();
        }
        mediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        mediaCodec.start();
    }

    private void StopEncoder() {
        try {
            mediaCodec.stop();
            mediaCodec.release();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void StopThread() {
        StopEncoder();
    }

    /**
     * MediaCodec只支持输入NV12格式数据
     */
    public void StartEncoderThread(byte[] input) {
        executor.execute(() -> {
            long pts = 0;
            long generateIndex = 0;

            if (input != null) {
                try {
                    ByteBuffer[] inputBuffers = mediaCodec.getInputBuffers();
                    ByteBuffer[] outputBuffers = mediaCodec.getOutputBuffers();
                    int inputBufferIndex = mediaCodec.dequeueInputBuffer(-1);
                    if (inputBufferIndex >= 0) {
                        pts = computePresentationTime(generateIndex);
                        ByteBuffer inputBuffer = inputBuffers[inputBufferIndex];
                        inputBuffer.clear();
                        inputBuffer.put(input);
                        mediaCodec.queueInputBuffer(inputBufferIndex, 0, input.length, pts, 0);
                        generateIndex += 1;
                    }

                    MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
                    int outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC);
                    while (outputBufferIndex >= 0) {
                        //Log.i("AvcEncoder", "Get H264 Buffer Success! flag = "+bufferInfo.flags+",pts = "+bufferInfo.presentationTimeUs+"");
                        ByteBuffer outputBuffer = outputBuffers[outputBufferIndex];
                        byte[] outData = new byte[bufferInfo.size];
                        outputBuffer.get(outData);
                        if (bufferInfo.flags == 2) {
                            configbyte = new byte[bufferInfo.size];
                            configbyte = outData;
                        } else if (bufferInfo.flags == 1) {
                            byte[] keyframe = new byte[bufferInfo.size + configbyte.length];
                            System.arraycopy(configbyte, 0, keyframe, 0, configbyte.length);
                            System.arraycopy(outData, 0, keyframe, configbyte.length, outData.length);
                            if (listener != null) {
                                listener.onResult(keyframe);
                            }
                        } else {
                            if (listener != null) {
                                listener.onResult(outData);
                            }
                        }

                        mediaCodec.releaseOutputBuffer(outputBufferIndex, false);
                        outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC);
                    }

                } catch (Throwable t) {
                    t.printStackTrace();
                }
            }
        });

    }

    /**
     * Generates the presentation time for frame N, in microseconds.
     */
    private long computePresentationTime(long frameIndex) {
        return 132 + frameIndex * 1000000 / m_framerate;
    }

    public interface OnListener {
        void onResult(byte[] outData);
    }
}